import os
import glob
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
import numpy as np

"""
the K-means effect of different anchors number 
"""

# 设置数据集标注文件所在的目录
dataset_directory = "/home/hw/dataset/cone_total_new/labels"

# 收集所有.txt标注文件
files = glob.glob(os.path.join(dataset_directory, "*.txt"))

# 初始化存储所有边界框的长和宽的列表
widths = []
heights = []

# 遍历所有文件
for file in files:
    # 排除classes.txt文件
    if "classes.txt" in file:
        continue

    with open(file, "r") as f:
        for line in f:
            parts = line.strip().split()
            # 确保数据行是有效的
            if len(parts) == 5:
                # 提取宽度和高度（归一化的，相对于图片的尺寸）
                _, _, w, h = map(float, parts[1:])
                widths.append(w)
                heights.append(h)

# 准备数据点为聚类
data = np.array(list(zip(widths, heights)))

# 指定聚类数量
num_clusters = 12  # 例如，分成15个聚类

# 创建KMeans实例，设置迭代次数为30
kmeans = KMeans(n_clusters=num_clusters, n_init=1, max_iter=300, random_state=42)
# 拟合模型
kmeans.fit(data)

# 获取聚类标签
labels = kmeans.labels_

# 获取聚类中心
centers = kmeans.cluster_centers_

# 绘制所有目标的长宽分布图，并按聚类标签着色
plt.figure(figsize=(10, 6))
for i in range(num_clusters):
    cluster_data = data[labels == i]
    # 散点图展示聚类数据
    plt.scatter(
        cluster_data[:, 0],
        cluster_data[:, 1],
        s=12,
        alpha=0.5,
        label=f"Cluster {i + 1}",
    )
    # 标记聚类中心
    plt.plot(centers[i, 0], centers[i, 1], "kx", markersize=6, markeredgewidth=2)

plt.title("K-means Clustering of Bounding Box Dimensions", fontsize=16)
plt.xlabel("Width (normalized)", fontsize=14)
plt.ylabel("Height (normalized)", fontsize=14)
plt.grid(True, linestyle="--", alpha=0.6)
plt.tick_params(axis="both", which="major", labelsize=12)
plt.legend(fontsize=12, loc="upper right")
plt.tight_layout()

# 保存图像为高分辨率的png文件
plt.savefig("kmeans_cluster_result.png", dpi=1000)

# 显示图像
plt.show()
